gpt-oss: Add MXFP8 x MXFP4 CUTLASS MOE for SM100 and BF16 x MXFP4 CUTLASS for SM90 + SwigluBias Activation #1396

djmmoss · 2025-08-06T04:26:30Z

📌 Description

This PR adds MXFP8 x MXFP4 CUTLASS MOE with SwigluBias for SM100 GPUs.

It also adds BF16 x MXFP4 CUTLASS MOE with SwigluBias for SM 90

🔍 Related Issues

N/A

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

I have installed pre-commit by running pip install pre-commit (or used your preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

Tests have been added or updated as needed.
All tests are passing (unittest, etc.).

Reviewer Notes

gemini-code-assist

Summary of Changes

Hello @djmmoss, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces support for MXFP8 x MXFP4 CUTLASS Mixture of Experts (MoE) operations, specifically targeting SM100 GPUs. I've also added a new SwigluBias activation function, which can be configured with per-expert parameters. These changes involve significant updates to the quantization pipeline, activation kernels, and the GEMM profiler backend to ensure efficient and accurate mixed-precision computations.

Highlights

New Mixed-Precision MoE Support: I've added comprehensive support for MXFP8 x MXFP4 CUTLASS Mixture of Experts (MoE) operations, specifically optimized for SM100 GPUs.
Introduction of SwigluBias Activation: A new SwigluBias activation function has been introduced, allowing for per-expert alpha, beta, and limit parameters to fine-tune activation behavior.
Enhanced Quantization Functions: The underlying quantization logic and kernels have been updated to seamlessly handle the new MXFP8 and MXFP4 data types, including refactoring of cvt_quant_get_sf_out_offset and quantize_with_block_size.
Improved Groupwise Scaling for WFP4A16: I've improved the handling of WFP4A16 groupwise scaling, with precise adjustments to stride and pointer calculations within the kernels.
Refactored Activation Kernels: The activation kernels (doGatedActivationKernel, doActivationKernel) have been refactored to use generic ActFn adaptors, making them more flexible and enabling the integration of SwigluBiasAdaptor.
Profiler Backend Updates: The GemmProfilerBackend has been updated to correctly account for WFP4A16 quantization and the new SwigluBias parameters when calculating workspace sizes.
Updated PyTorch Bindings: The PyTorch bindings (flashinfer_cutlass_fused_moe_sm100_ops.cu) have been modified to pass the new ActivationParams struct and handle swiglu_alpha, swiglu_beta, and swiglu_limit tensors from Python.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request adds support for MXFP8 x MXFP4 CUTLASS MoE with SwigluBias activation for SM100 GPUs. The changes are extensive, including new template instantiations, refactoring of quantization kernels for better generality, and the addition of activation adaptors to handle the new SwigluBias activation. The code is well-structured, and the new features seem to be implemented correctly. I have one suggestion to improve maintainability by refactoring a piece of duplicated code.

gemini-code-assist · 2025-08-06T04:30:09Z

csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh

+    float gate_alpha = 1.0f;
+    float gate_beta = 0.0f;
+    float gate_limit = std::numeric_limits<float>::infinity();
+    if (bias_ptr || IsNVFP4 || IsMXFP8 || use_per_expert_act_scale ||
+        activation_params.swiglu_alpha || activation_params.swiglu_beta ||
+        activation_params.swiglu_limit) {
      // TODO this is almost certainly faster as a linear scan
      expert =
          findTotalEltsLessThanTarget(expert_first_token_offset, num_experts_per_node, token + 1) -
          1;
+      gate_alpha = activation_params.swiglu_alpha ? activation_params.swiglu_alpha[expert] : 1.0f;
+      gate_beta = activation_params.swiglu_beta ? activation_params.swiglu_beta[expert] : 0.0f;
+      gate_limit = activation_params.swiglu_limit ? activation_params.swiglu_limit[expert]
+                                                  : std::numeric_limits<float>::infinity();
    }


The logic for extracting SwiGLU parameters (gate_alpha, gate_beta, gate_limit) is very similar to the logic in doGatedActivationKernel at lines 2068-2079. To improve maintainability and reduce code duplication, consider refactoring this into a helper function. This function could take activation_params and the expert index as input and return a struct or tuple with the gate parameters.

aleozlx · 2025-08-13T22:58:13Z

flashinfer/__init__.py

    e2m1_and_ufp8sf_scale_to_float,
    fp4_quantize,
-    nvfp4_block_scale_interleave,


perhaps non-blocking comment, just an alert:

renaming this may break cached weight processing (sgl will use the same code as the tests)
https://github.com/flashinfer-ai/flashinfer/pull/1412/files#diff-d2eee7647631902c9b37dc7b03c56129e3f5b94f88df67cdf18edeff5c03d159R421

makes sense, I've add this to handle that: https://github.com/flashinfer-ai/flashinfer/pull/1396/files#diff-f304f5c1338c89b675fb67a992127d1c4ba5164a7b2a0dc874420f790bdde0f6R355

yzh119 · 2025-08-14T00:03:12Z

flashinfer/fp4_quantization.py

+    if get_device_arch() == "100a":
+        nvcc_flags += sm100a_nvcc_flags
+    elif get_device_arch() == "90a":
+        nvcc_flags += sm90a_nvcc_flags


Does sm90a support native fp4 cvt ptx instructions?

it doesn't I've reverted back some of the changes to reflect that

yzh119 · 2025-08-14T00:04:05Z

flashinfer/fp4_quantization.py

@@ -299,7 +345,7 @@ def fp4_quantize(
    return x_q, sf


-def nvfp4_block_scale_interleave(unswizzled_sf: torch.Tensor) -> torch.Tensor:
+def block_scale_interleave(unswizzled_sf: torch.Tensor) -> torch.Tensor:


I would suggest also keeping nvfp4_block_scale_interleave for backward compatibility, as mentioned by @aleozlx

nvfp4_block_scale_interleave = block_scale_interleave

see: #1396 (comment)

yzh119 · 2025-08-14T00:07:09Z

include/flashinfer/comm/trtllm_allreduce_fusion.cuh

@@ -16,7 +16,7 @@ namespace flashinfer {

 namespace trtllm_allreduce_fusion {

-enum class FP4QuantizationSFLayout {
+enum class QuantizationSFLayout {


Can we keep a single source of QuantizationSFLayout? This enum class have been defined several times in different places, I would recommend a standalone header such as include/flashinfer/quantization.cuh for hosting these common classes.

there is a larger refactor of the quantization utilities through the entire repo that can be done. Can this be handled in a separate PR along with those changes?

flashinfer/fused_moe/core.py

joker-eph · 2025-08-14T14:33:58Z

flashinfer/fused_moe/core.py

@@ -629,6 +605,9 @@ def cutlass_fused_moe(
    fc1_expert_biases: Optional[torch.Tensor] = None,
    fc2_expert_biases: Optional[torch.Tensor] = None,
    input_sf: Optional[torch.Tensor] = None,
+    swiglu_alpha: Optional[torch.Tensor] = None,
+    swiglu_beta: Optional[torch.Tensor] = None,
+    swiglu_limit: Optional[torch.Tensor] = None,


Can you update the docstring to document the newly added arguments?

done, please check

joker-eph · 2025-08-14T14:36:38Z

flashinfer/fp4_quantization.py

+        nvcc_flags += sm100a_nvcc_flags
+    elif get_device_arch() == "90a":
+        nvcc_flags += sm90a_nvcc_flags
+    else:


Isn't this an issue to query the current GPU here? Wouldn't this break AOT build where we want to build independently of even having a GPU available?

makes sense, I've split it out and created two gen functions which are now included in the aot please let me know if that works.

flashinfer/__init__.py

joker-eph · 2025-08-14T14:47:49Z

flashinfer/utils.py

+@functools.cache
+def get_device_arch():
+    major, minor = torch.cuda.get_device_capability()
+    suffix = "a" if major >= 9 else ""


Can we support sm100f instead of sm100a to support the family out-of-the-box?

I think I've removed this issue now

flashinfer/jit/core.py

joker-eph

LG, I'll let @yzh119 have another look (tomorrow probably)

djmmoss · 2025-08-15T18:48:51Z

Hold off for now, I've run into a small error with the AOT build, this cleans up kernels that should be part of the build...

djmmoss added 3 commits August 5, 2025 15:24

add mxfp8 x mxfp4 cutlass fused moe

5a9f697

add swiglu test

6d2ccc2

pre-commit checks

542ca4b

gemini-code-assist bot reviewed Aug 6, 2025

View reviewed changes

djmmoss changed the title ~~gpt-oss: Add MXFP8 x MXFP4 CUTLASS MOE + SwigluBias Activation for SM100~~ gpt-oss: Add MXFP8 x MXFP4 CUTLASS MOE for SM100 and BF16 x MXFP4 CUTLASS for SM90 + SwigluBias Activation Aug 6, 2025

djmmoss added 5 commits August 8, 2025 09:59

Merge remote-tracking branch 'origin/main' into dmoss/cutlass-mxfp8-moe

8ad7c23

wip

e13b3b2

Merge remote-tracking branch 'origin/main' into dmoss/cutlass-mxfp8-moe

998c39f

final changes to enable bf16xmxfp4 moe for hopper

fb6b9e2

Merge remote-tracking branch 'origin/main' into dmoss/cutlass-mxfp8-moe

ecd4133

aleozlx reviewed Aug 13, 2025

View reviewed changes

yzh119 reviewed Aug 14, 2025

View reviewed changes

joker-eph reviewed Aug 14, 2025

View reviewed changes

flashinfer/fused_moe/core.py Show resolved Hide resolved

joker-eph reviewed Aug 14, 2025

View reviewed changes

flashinfer/__init__.py Show resolved Hide resolved

joker-eph reviewed Aug 14, 2025

View reviewed changes

djmmoss added 4 commits August 14, 2025 10:03

addressing PR comments

f6f01c4

address pr comments

aad97d5

minor

e8d94a8

typo

5949771

joker-eph reviewed Aug 14, 2025

View reviewed changes

flashinfer/jit/core.py Show resolved Hide resolved

device_arch updates

da99130

joker-eph approved these changes Aug 14, 2025

View reviewed changes

gpt-oss: Add MXFP8 x MXFP4 CUTLASS MOE for SM100 and BF16 x MXFP4 CUTLASS for SM90 + SwigluBias Activation #1396

Are you sure you want to change the base?

gpt-oss: Add MXFP8 x MXFP4 CUTLASS MOE for SM100 and BF16 x MXFP4 CUTLASS for SM90 + SwigluBias Activation #1396

Uh oh!

Conversation

djmmoss commented Aug 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

djmmoss Aug 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

joker-eph left a comment

Choose a reason for hiding this comment

Uh oh!

djmmoss commented Aug 15, 2025

Uh oh!

Uh oh!

djmmoss commented Aug 6, 2025 •

edited

Loading

djmmoss Aug 14, 2025 •

edited

Loading